#include <iostream>
#include <vector>
#include <ctime>
#include <numeric>

#include "tensor.h"
#include "covariance.h"
#include "mtrand.h"
#include "dataset.h"
#include "pdf.h"
#include <vector>
#include <sstream>

#include "dataset.h"
#include "pdf.h"

#include "TFile.h"
#include "TChain.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TProfile2D.h"
#include "TF1.h"
#include "TF2.h"
#include "TF3.h"
#include "TLorentzVector.h"
#include "TGraphErrors.h"
#include "TGraph2DErrors.h"
#include "make_list.h"

using namespace std;

//initialize random number generator
MTRand_int32 irand(time(0));

const double PI=3.14159265359;


class Vec
{
public:
  vector<double> x,y,err;

  Vec(){}

  void push_back(double x_, double y_, double yerr_)
  {
    x.push_back(x_);
    y.push_back(y_);
    err.push_back(yerr_);    
  }

  TGraphErrors get_graph(string name, bool logx=false) const
  {
    vector<double> xprime(x);
    if(logx)
      {
	for(unsigned int i=0; i<xprime.size(); ++i)
	  {
	    xprime[i]=exp10(-2*xprime[i]);
	  }
      }
    TGraphErrors result(xprime.size(), &xprime[0], &y[0], NULL, &err[0]);
    result.SetTitle(name.c_str());
    result.SetName(name.c_str());

    return result;
  }

  //get a 1D graph from pt
  void fill_data(PDF& pdf, double pt)
  {
    PDF::MC_Generator gen=
      pdf.Generator(make_list<PDF::MC_Flag>
		    (PDF::OUTPUT, PDF::INPUT));

    vector<PDF::MC_Event> evt=
      gen.GenerateFull(make_list<double>(log(pt/50.)));

    for(unsigned int i=0; i<evt.size(); ++i)
      {
	x.push_back(evt[i].value[0]);
	y.push_back(evt[i].weight);
	err.push_back(evt[i].err);
      }
  }

};



float DeltaR(float eta, float phi)
{
  if(fabs(phi)> PI)
    phi= PI-fabs(phi);
  return sqrt(eta*eta + phi*phi);
}

int main(int argc, char **argv)
{
  if(argc < 3+1)
    {
      cout<<"Usage: ./sherpa_PDF_output sherpa_input.root madgraph_input_gg madgraph_input_pp"
	  <<endl;
      throw;
    }

  cout<<"madgraph inputs are read as madgraph_input.?.root ranges from 1 to 9"
      <<endl;
  cout<<"Outputs are automatically stored in PDF_*"<<endl;
  
  //open rootfile
  //for sherpa files

  TChain chain("Ntuple");
  chain.Add(argv[1]);

  //plots of tau21 and jet mass
  map<string, TH1D*> hist;
  const int pt_bins=20;
  const int pt_width=50;
  
  for(unsigned int tpt=50; tpt<=1000; tpt+=50)
    {
      stringstream sout;
      sout<<"gg_t21_1_"<<tpt;
      hist[sout.str()]=new TH1D(sout.str().c_str(),
				sout.str().c_str(),
				20, 0, 1.0);

      sout.str("");
      sout<<"gg_t21_2_"<<tpt;
      hist[sout.str()]=new TH1D(sout.str().c_str(),
				sout.str().c_str(),
				20, 0, 1.0);


      sout.str("");
      sout<<"gg_mpt_1_"<<tpt;
      hist[sout.str()]=new TH1D(sout.str().c_str(),
				sout.str().c_str(),
				20, 0, 3.0);

      sout.str("");
      sout<<"gg_mpt_2_"<<tpt;
      hist[sout.str()]=new TH1D(sout.str().c_str(),
				sout.str().c_str(),
				20, 0, 3.0);
 
      sout.str("");
      sout<<"qq_t21_1_"<<tpt;
      hist[sout.str()]=new TH1D(sout.str().c_str(),
				sout.str().c_str(),
				20, 0, 1.0);

      sout.str("");
      sout<<"qq_t21_2_"<<tpt;
      hist[sout.str()]=new TH1D(sout.str().c_str(),
				sout.str().c_str(),
				20, 0, 1.0);


      sout.str("");
      sout<<"qq_mpt_1_"<<tpt;
      hist[sout.str()]=new TH1D(sout.str().c_str(),
				sout.str().c_str(),
				20, 0, 3.0);

      sout.str("");
      sout<<"qq_mpt_2_"<<tpt;
      hist[sout.str()]=new TH1D(sout.str().c_str(),
				sout.str().c_str(),
				20, 0, 3.0);

    }


  for(map<string, TH1D*>::iterator it=hist.begin();
      it!=hist.end(); ++it)
    {
      it->second->Sumw2();
    }


  //declare variables
  const int maxsize=20;
  Int_t           size;
  Double_t         w;   //[size]
  Float_t         pt[maxsize];   //[size]
  Float_t         eta[maxsize];
  Float_t         phi[maxsize];
  Float_t         m[maxsize];   //[size]
  Float_t         t1[maxsize];
  Float_t         t2[maxsize];
  Float_t         t3[maxsize];


  //read the input file
  chain.SetBranchAddress("size", &size);
  chain.SetBranchAddress("w", &w);
  chain.SetBranchAddress("pt", pt);
  chain.SetBranchAddress("eta", eta);
  chain.SetBranchAddress("phi", phi);
  chain.SetBranchAddress("m", m);
  chain.SetBranchAddress("tau1", t1);
  chain.SetBranchAddress("tau2", t2);
  chain.SetBranchAddress("tau3", t3);


  DataSet DataSet_mpt_pt(make_list<double>(0, 0),
			 make_list<double>(3.0, 5.0),
			 make_list<double>(0.03, 0.05));
  
  DataSet DataSet_tau21_pt(make_list<double>(0, 0),
			   make_list<double>(1.0, 5.0),
			   make_list<double>(0.02, 0.05));


  //list of DataSet
  vector<DataSet*> dataset_mpt_njet(4);
  vector<DataSet*> dataset_tau21_njet(4);
  
  for(int n=0; n<4; n++)
    {
      dataset_mpt_njet[n]=new
	DataSet(make_list<double>(0, 0),
		make_list<double>(3.0, 5.0),
		make_list<double>(0.03, 0.05));
  
      dataset_tau21_njet[n]=new
	DataSet(make_list<double>(0, 0),
		make_list<double>(1.0, 5.0),
		make_list<double>(0.02, 0.05));
    }

  
  /////
  DataSet DataSet_mpt_pt_weight(make_list<double>(0, 0),
				make_list<double>(3.0, 5.0),
				make_list<double>(0.03, 0.05));
  
  DataSet DataSet_tau21_pt_weight(make_list<double>(0, 0),
				  make_list<double>(1.0, 5.0),
				  make_list<double>(0.02, 0.05));

  DataSet DataSet_mpt_pt_gg(make_list<double>(0, 0),
			    make_list<double>(3.0, 5.0),
			    make_list<double>(0.03, 0.05));
  
  DataSet DataSet_tau21_pt_gg(make_list<double>(0, 0),
			      make_list<double>(1.0, 5.0),
			      make_list<double>(0.02, 0.05));


  DataSet DataSet_mpt_pt_qq(make_list<double>(0, 0),
			    make_list<double>(3.0, 5.0),
			    make_list<double>(0.03, 0.05));
  
  DataSet DataSet_tau21_pt_qq(make_list<double>(0, 0),
			      make_list<double>(1.0, 5.0),
			      make_list<double>(0.02, 0.05));
  
  
  
  //loop over events
  double total_weight=0;
  int nevt=0;
  for(Long64_t i=0; i<chain.GetEntries(); i++)
    {
      
      if(i%10000==0)
	cout<<"\r\033[K"<<(i*10000/chain.GetEntries())*0.01<<"% trained..."
	    <<flush;

      chain.GetEntry(i);
      if(size<2) continue;
      
      //some sanity cuts
      if(t1[0] < 1e-5) continue;
      if(t1[1] < 1e-5) continue;
      if(t2[0] < 1e-5) continue;
      if(t2[1] < 1e-5) continue;
      if(pt[0] < 50.) continue;
      if(pt[1] < 50.) continue;
      if(m[0]<20) continue;
      if(m[1]<20) continue;

      
      for(int n=0; n<size && 
	    n<dataset_mpt_njet.size() && 
	    n<dataset_tau21_njet.size(); ++n)
	{
	  if(t1[n] < 1e-5) continue;
	  if(t2[n] < 1e-5) continue;
	  if(pt[n] < 50.) continue;
	  if(m[n]<20) continue;

	  double logpt=log(pt[n]/50.);
	  double mpt=-log10(m[n]/pt[n]);
	  double t21=t2[n]/t1[n];

	  dataset_mpt_njet[n]->Fill(make_list<double>(mpt,logpt));
	  dataset_tau21_njet[n]->Fill(make_list<double>(t21,logpt));
	}

      ++nevt;

      double dR=DeltaR(eta[0]-eta[1], phi[0]-phi[1]);
      
      TLorentzVector j1,j2;
      j1.SetPtEtaPhiM(pt[0],eta[0],phi[0],m[0]);
      j2.SetPtEtaPhiM(pt[1],eta[1],phi[1],m[1]);
      double mjj=(j1+j2).M();
      double logmjj=log(mjj/50.);

      double logpt1=log(pt[0]/50.);
      double logpt2=log(pt[1]/50.);
      /*
      double mpt1=m[0]/pt[0];
      double mpt2=m[1]/pt[1];
      */
      double mpt1=-log10(m[0]/pt[0]);
      double mpt2=-log10(m[1]/pt[1]);

      double t21_1=t2[0]/t1[0];
      double t21_2=t2[1]/t1[1];


      DataSet_tau21_pt.Fill(make_list<double>(t21_1, logpt1));
      DataSet_tau21_pt.Fill(make_list<double>(t21_2, logpt2));

      DataSet_mpt_pt.Fill(make_list<double>(mpt1,logpt1 ));
      DataSet_mpt_pt.Fill(make_list<double>(mpt2,logpt2 ));


      DataSet_tau21_pt_weight.Fill(make_list<double>(t21_1, logpt1),w);
      DataSet_tau21_pt_weight.Fill(make_list<double>(t21_2, logpt2),w);

      DataSet_mpt_pt_weight.Fill(make_list<double>(mpt1,logpt1 ),w);
      DataSet_mpt_pt_weight.Fill(make_list<double>(mpt2,logpt2 ),w);
      
      total_weight+=w;
    }

  cout<<endl;

  //here compute come PDF
  PDF PDF_tau21_pt=DataSet_tau21_pt.ComputePDF(true);
  PDF PDF_mpt_pt=DataSet_mpt_pt.ComputePDF(true);

  PDF PDF_tau21_pt_weight=DataSet_tau21_pt_weight.ComputePDF(true);
  PDF PDF_mpt_pt_weight=DataSet_mpt_pt_weight.ComputePDF(true);

  //get PDF list
  PDF PDF_mpt_njet[4];
  PDF PDF_tau21_njet[4]; 

  //for the njet cases
  for(int n=0; n<dataset_mpt_njet.size(); ++n)
    {
      PDF_mpt_njet[n]=dataset_mpt_njet[n]->ComputePDF();
      PDF_tau21_njet[n]=dataset_tau21_njet[n]->ComputePDF();      
    }

  //now process madgraph events
  for(unsigned int k=1; k<=12; k++)
    {
      TChain chain_madgraph("Ntuple");
      stringstream sout;
      sout<<argv[2]<<"."<<k<<".root";
      chain_madgraph.Add(sout.str().c_str());

      //read the input file
      chain_madgraph.SetBranchAddress("size", &size);
      chain_madgraph.SetBranchAddress("w", &w);
      chain_madgraph.SetBranchAddress("pt", pt);
      chain_madgraph.SetBranchAddress("eta", eta);
      chain_madgraph.SetBranchAddress("phi", phi);
      chain_madgraph.SetBranchAddress("m", m);
      chain_madgraph.SetBranchAddress("tau1", t1);
      chain_madgraph.SetBranchAddress("tau2", t2);
      chain_madgraph.SetBranchAddress("tau3", t3);

      double rescale=1.0/chain_madgraph.GetEntries();

      for(unsigned int i=0; i<chain_madgraph.GetEntries(); ++i)
	{
	  if(i%1000==0)
	    cout<<"\r\033[K"<<(i*10000/chain_madgraph.GetEntries())*0.01
		<<"% trained..."
		<<flush;
	  
	  
	  chain_madgraph.GetEntry(i);
	  if(size<2) continue;
	  
	  //some sanity cuts
	  if(t1[0] < 1e-5) continue;
	  if(t1[1] < 1e-5) continue;
	  if(t2[0] < 1e-5) continue;
	  if(t2[1] < 1e-5) continue;
	  if(pt[0] < 50.) continue;
	  if(pt[1] < 50.) continue;
	  if(m[0]<20) continue;
	  if(m[1]<20) continue;
	  
	  double dR=DeltaR(eta[0]-eta[1], phi[0]-phi[1]);
	  
	  TLorentzVector j1,j2;
	  j1.SetPtEtaPhiM(pt[0],eta[0],phi[0],m[0]);
	  j2.SetPtEtaPhiM(pt[1],eta[1],phi[1],m[1]);
	  double mjj=(j1+j2).M();
	  double logmjj=log(mjj/50.);
	  
	  double logpt1=log(pt[0]/50.);
	  double logpt2=log(pt[1]/50.);
	  /*
	    double mpt1=m[0]/pt[0];
	    double mpt2=m[1]/pt[1];
	  */
	  double mpt1=-log10(m[0]/pt[0]);
	  double mpt2=-log10(m[1]/pt[1]);
	  
	  double t21_1=t2[0]/t1[0];
	  double t21_2=t2[1]/t1[1];

	  DataSet_tau21_pt_gg.Fill(make_list<double>(t21_1, logpt1), 
				   w*rescale);
	  DataSet_tau21_pt_gg.Fill(make_list<double>(t21_2, logpt2),
				   w*rescale);
	  
	  DataSet_mpt_pt_gg.Fill(make_list<double>(mpt1,logpt1 ),
				 w*rescale);
	  DataSet_mpt_pt_gg.Fill(make_list<double>(mpt2,logpt2 ),
				 w*rescale);


	  for(unsigned int tpt=50; tpt<=1000; tpt+=50)
	    if(pt[0] < tpt+50 || tpt==1000)
	      {
		stringstream sout;
		sout<<"gg_t21_1_"<<tpt;
		hist[sout.str()]->Fill(t21_1, w*rescale);
		
		sout.str("");
		sout<<"gg_mpt_1_"<<tpt;
		hist[sout.str()]->Fill(mpt1, w*rescale);

		break;
	      }

	  for(unsigned int tpt=50; tpt<=1000; tpt+=50)
	    if(pt[1] < tpt+50 || tpt==1000)
	      {
		stringstream sout;
		sout<<"gg_t21_2_"<<tpt;
		hist[sout.str()]->Fill(t21_2, w*rescale);
		
		sout.str("");
		sout<<"gg_mpt_2_"<<tpt;
		hist[sout.str()]->Fill(mpt2, w*rescale);

		break;
	      }

	  
	}
      
      cout<<"gg "<<k<<" file processed."<<endl;
    }
  //here compute come PDF
  PDF PDF_tau21_pt_gg=DataSet_tau21_pt_gg.ComputePDF(true);
  PDF PDF_mpt_pt_gg=DataSet_mpt_pt_gg.ComputePDF(true);
  
  
  //now process qq evnts
  for(unsigned int k=1; k<=12; k++)
    {
      TChain chain_madgraph("Ntuple");
      stringstream sout;
      sout<<argv[3]<<"."<<k<<".root";
      chain_madgraph.Add(sout.str().c_str());

      //read the input file
      chain_madgraph.SetBranchAddress("size", &size);
      chain_madgraph.SetBranchAddress("w", &w);
      chain_madgraph.SetBranchAddress("pt", pt);
      chain_madgraph.SetBranchAddress("eta", eta);
      chain_madgraph.SetBranchAddress("phi", phi);
      chain_madgraph.SetBranchAddress("m", m);
      chain_madgraph.SetBranchAddress("tau1", t1);
      chain_madgraph.SetBranchAddress("tau2", t2);
      chain_madgraph.SetBranchAddress("tau3", t3);

      double rescale=1.0/chain_madgraph.GetEntries();

      for(unsigned int i=0; i<chain_madgraph.GetEntries(); ++i)
	{
	  if(i%1000==0)
	    cout<<"\r\033[K"<<(i*10000/chain_madgraph.GetEntries())*0.01
		<<"% trained..."
		<<flush;
	  
	  
	  chain_madgraph.GetEntry(i);
	  if(size<2) continue;
	  
	  //some sanity cuts
	  if(t1[0] < 1e-5) continue;
	  if(t1[1] < 1e-5) continue;
	  if(t2[0] < 1e-5) continue;
	  if(t2[1] < 1e-5) continue;
	  if(pt[0] < 50.) continue;
	  if(pt[1] < 50.) continue;
	  if(m[0]<20) continue;
	  if(m[1]<20) continue;
	  
	  double dR=DeltaR(eta[0]-eta[1], phi[0]-phi[1]);
	  
	  TLorentzVector j1,j2;
	  j1.SetPtEtaPhiM(pt[0],eta[0],phi[0],m[0]);
	  j2.SetPtEtaPhiM(pt[1],eta[1],phi[1],m[1]);
	  double mjj=(j1+j2).M();
	  double logmjj=log(mjj/50.);
	  
	  double logpt1=log(pt[0]/50.);
	  double logpt2=log(pt[1]/50.);
	  /*
	    double mpt1=m[0]/pt[0];
	    double mpt2=m[1]/pt[1];
	  */
	  double mpt1=-log10(m[0]/pt[0]);
	  double mpt2=-log10(m[1]/pt[1]);
	  
	  double t21_1=t2[0]/t1[0];
	  double t21_2=t2[1]/t1[1];

	  DataSet_tau21_pt_qq.Fill(make_list<double>(t21_1, logpt1), 
				   w*rescale);
	  DataSet_tau21_pt_qq.Fill(make_list<double>(t21_2, logpt2),
				   w*rescale);
	  
	  DataSet_mpt_pt_qq.Fill(make_list<double>(mpt1,logpt1 ),
				 w*rescale);
	  DataSet_mpt_pt_qq.Fill(make_list<double>(mpt2,logpt2 ),
				 w*rescale);


	  for(unsigned int tpt=50; tpt<=1000; tpt+=50)
	    if(pt[0] < tpt+50 || tpt==1000)
	      {
		stringstream sout;
		sout<<"qq_t21_1_"<<tpt;
		hist[sout.str()]->Fill(t21_1, w*rescale);
		
		sout.str("");
		sout<<"qq_mpt_1_"<<tpt;
		hist[sout.str()]->Fill(mpt1, w*rescale);

		break;
	      }

	  for(unsigned int tpt=50; tpt<=1000; tpt+=50)
	    if(pt[1] < tpt+50 || tpt==1000)
	      {
		stringstream sout;
		sout<<"qq_t21_2_"<<tpt;
		hist[sout.str()]->Fill(t21_2, w*rescale);
		
		sout.str("");
		sout<<"qq_mpt_2_"<<tpt;
		hist[sout.str()]->Fill(mpt2, w*rescale);

		break;
	      }

	  
	}
      
      cout<<"qq "<<k<<" file processed."<<endl;
    }
  //here compute come PDF
  PDF PDF_tau21_pt_qq=DataSet_tau21_pt_qq.ComputePDF(true);
  PDF PDF_mpt_pt_qq=DataSet_mpt_pt_qq.ComputePDF(true);


  //get output
  TFile out("output.root", "RECREATE");
  map<string, Vec> graph_list;

  for(int i=60; i<=3000; i+=10)
    {
      stringstream sout;
      
      sout.str("");
      sout<<"t21_sherpa_"<<i;
      graph_list[sout.str()].fill_data(PDF_tau21_pt, i);
      graph_list[sout.str()].get_graph(sout.str()).Write();

      sout.str("");
      sout<<"t21_sherpa_weight_"<<i;
      graph_list[sout.str()].fill_data(PDF_tau21_pt_weight, i);
      graph_list[sout.str()].get_graph(sout.str()).Write();

      sout.str("");
      sout<<"t21_gg_"<<i;
      graph_list[sout.str()].fill_data(PDF_tau21_pt_gg, i);
      graph_list[sout.str()].get_graph(sout.str()).Write();

      sout.str("");
      sout<<"t21_qq_"<<i;
      graph_list[sout.str()].fill_data(PDF_tau21_pt_qq, i);
      graph_list[sout.str()].get_graph(sout.str()).Write();


      sout.str("");
      sout<<"mpt_sherpa_"<<i;
      graph_list[sout.str()].fill_data(PDF_mpt_pt, i);
      graph_list[sout.str()].get_graph(sout.str(),true).Write();

      sout.str("");
      sout<<"mpt_sherpa_weight_"<<i;
      graph_list[sout.str()].fill_data(PDF_mpt_pt_weight, i);
      graph_list[sout.str()].get_graph(sout.str(),true).Write();

      sout.str("");
      sout<<"mpt_gg_"<<i;
      graph_list[sout.str()].fill_data(PDF_mpt_pt_gg, i);
      graph_list[sout.str()].get_graph(sout.str(),true).Write();

      sout.str("");
      sout<<"mpt_qq_"<<i;
      graph_list[sout.str()].fill_data(PDF_mpt_pt_qq, i);
      graph_list[sout.str()].get_graph(sout.str(),true).Write();


      //do the njet case
      for(int n=0; n<4; ++n)
	{
	  sout.str("");
	  sout<<"t21_sherpa_"<<n<<"jet_"<<i;
	  graph_list[sout.str()].fill_data(PDF_tau21_njet[n], i);
	  graph_list[sout.str()].get_graph(sout.str()).Write();

	  sout.str("");
	  sout<<"mpt_sherpa_"<<n<<"jet_"<<i;
	  graph_list[sout.str()].fill_data(PDF_mpt_njet[n], i);
	  graph_list[sout.str()].get_graph(sout.str(), true).Write();	  
	}

    }


  //write histograms
  for(map<string, TH1D*>::iterator it=hist.begin();
      it!=hist.end(); ++it)
    {
      it->second->Write();
    }

  cout<<"Program ends."<<endl;
  return 1;
}
  
